import java.util.Arrays;

public class Main {

    // S盒
    private static final int[] S_BOX = {
            0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5, 0x30, 0x01, 0x67, 0x2B, 0xFE, 0xD7, 0xAB, 0x76,
            0xCA, 0x82, 0xC9, 0x7D, 0xFA, 0x59, 0x47, 0xF0, 0xAD, 0xD4, 0xA2, 0xAF, 0x9C, 0xA4, 0x72, 0xC0,
            0xB7, 0xFD, 0x93, 0x26, 0x36, 0x3F, 0xF7, 0xCC, 0x34, 0xA5, 0xE5, 0xF1, 0x71, 0xD8, 0x31, 0x15,
            0x04, 0xC7, 0x23, 0xC3, 0x18, 0x96, 0x05, 0x9A, 0x07, 0x12, 0x80, 0xE2, 0xEB, 0x27, 0xB2, 0x75,
            0x09, 0x83, 0x2C, 0x1A, 0x1B, 0x6E, 0x5A, 0xA0, 0x52, 0x3B, 0xD6, 0xB3, 0x29, 0xE3, 0x2F, 0x84,
            0x53, 0xD1, 0x00, 0xED, 0x20, 0xFC, 0xB1, 0x5B, 0x6A, 0xCB, 0xBE, 0x39, 0x4A, 0x4C, 0x58, 0xCF,
            0xD0, 0xEF, 0xAA, 0xFB, 0x43, 0x4D, 0x33, 0x85, 0x45, 0xF9, 0x02, 0x7F, 0x50, 0x3C, 0x9F, 0xA8,
            0x51, 0xA3, 0x40, 0x8F, 0x92, 0x9D, 0x38, 0xF5, 0xBC, 0xB6, 0xDA, 0x21, 0x10, 0xFF, 0xF3, 0xD2,
            0xCD, 0x0C, 0x13, 0xEC, 0x5F, 0x97, 0x44, 0x17, 0xC4, 0xA7, 0x7E, 0x3D, 0x64, 0x5D, 0x19, 0x73,
            0x60, 0x81, 0x4F, 0xDC, 0x22, 0x2A, 0x90, 0x88, 0x46, 0xEE, 0xB8, 0x14, 0xDE, 0x5E, 0x0B, 0xDB,
            0xE0, 0x32, 0x3A, 0x0A, 0x49, 0x06, 0x24, 0x5C, 0xC2, 0xD3, 0xAC, 0x62, 0x91, 0x95, 0xE4, 0x79,
            0xE7, 0xC8, 0x37, 0x6D, 0x8D, 0xD5, 0x4E, 0xA9, 0x6C, 0x56, 0xF4, 0xEA, 0x65, 0x7A, 0xAE, 0x08,
            0xBA, 0x78, 0x25, 0x2E, 0x1C, 0xA6, 0xB4, 0xC6, 0xE8, 0xDD, 0x74, 0x1F, 0x4B, 0xBD, 0x8B, 0x8A,
            0x70, 0x3E, 0xB5, 0x66, 0x48, 0x03, 0xF6, 0x0E, 0x61, 0x35, 0x57, 0xB9, 0x86, 0xC1, 0x1D, 0x9E,
            0xE1, 0xF8, 0x98, 0x11, 0x69, 0xD9, 0x8E, 0x94, 0x9B, 0x1E, 0x87, 0xE9, 0xCE, 0x55, 0x28, 0xDF,
            0x8C, 0xA1, 0x89, 0x0D, 0xBF, 0xE6, 0x42, 0x68, 0x41, 0x99, 0x2D, 0x0F, 0xB0, 0x54, 0xBB, 0x16
    };
    // 逆S盒
    private static final int[] INV_S_BOX = {
            0x52, 0x09, 0x6A, 0xD5, 0x30, 0x36, 0xA5, 0x38, 0xBF, 0x40, 0xA3, 0x9E, 0x81, 0xF3, 0xD7, 0xFB,
            0x7C, 0xE3, 0x39, 0x82, 0x9B, 0x2F, 0xFF, 0x87, 0x34, 0x8E, 0x43, 0x44, 0xC4, 0xDE, 0xE9, 0xCB,
            0x54, 0x7B, 0x94, 0x32, 0xA6, 0xC2, 0x23, 0x3D, 0xEE, 0x4C, 0x95, 0x0B, 0x42, 0xFA, 0xC3, 0x4E,
            0x08, 0x2E, 0xA1, 0x66, 0x28, 0xD9, 0x24, 0xB2, 0x76, 0x5B, 0xA2, 0x49, 0x6D, 0x8B, 0xD1, 0x25,
            0x72, 0xF8, 0xF6, 0x64, 0x86, 0x68, 0x98, 0x16, 0xD4, 0xA4, 0x5C, 0xCC, 0x5D, 0x65, 0xB6, 0x92,
            0x6C, 0x70, 0x48, 0x50, 0xFD, 0xED, 0xB9, 0xDA, 0x5E, 0x15, 0x46, 0x57, 0xA7, 0x8D, 0x9D, 0x84,
            0x90, 0xD8, 0xAB, 0x00, 0x8C, 0xBC, 0xD3, 0x0A, 0xF7, 0xE4, 0x58, 0x05, 0xB8, 0xB3, 0x45, 0x06,
            0xD0, 0x2C, 0x1E, 0x8F, 0xCA, 0x3F, 0x0F, 0x02, 0xC1, 0xAF, 0xBD, 0x03, 0x01, 0x13, 0x8A, 0x6B,
            0x3A, 0x91, 0x11, 0x41, 0x4F, 0x67, 0xDC, 0xEA, 0x97, 0xF2, 0xCF, 0xCE, 0xF0, 0xB4, 0xE6, 0x73,
            0x96, 0xAC, 0x74, 0x22, 0xE7, 0xAD, 0x35, 0x85, 0xE2, 0xF9, 0x37, 0xE8, 0x1C, 0x75, 0xDF, 0x6E,
            0x47, 0xF1, 0x1A, 0x71, 0x1D, 0x29, 0xC5, 0x89, 0x6F, 0xB7, 0x62, 0x0E, 0xAA, 0x18, 0xBE, 0x1B,
            0xFC, 0x56, 0x3E, 0x4B, 0xC6, 0xD2, 0x79, 0x20, 0x9A, 0xDB, 0xC0, 0xFE, 0x78, 0xCD, 0x5A, 0xF4,
            0x1F, 0xDD, 0xA8, 0x33, 0x88, 0x07, 0xC7, 0x31, 0xB1, 0x12, 0x10, 0x59, 0x27, 0x80, 0xEC, 0x5F,
            0x60, 0x51, 0x7F, 0xA9, 0x19, 0xB5, 0x4A, 0x0D, 0x2D, 0xE5, 0x7A, 0x9F, 0x93, 0xC9, 0x9C, 0xEF,
            0xA0, 0xE0, 0x3B, 0x4D, 0xAE, 0x2A, 0xF5, 0xB0, 0xC8, 0xEB, 0xBB, 0x3C, 0x83, 0x53, 0x99, 0x61,
            0x17, 0x2B, 0x04, 0x7E, 0xBA, 0x77, 0xD6, 0x26, 0xE1, 0x69, 0x14, 0x63, 0x55, 0x21, 0x0C, 0x7D
    };

    // 轮常数
    private static final int[] RCON = { 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1B, 0x36 };

    // 明文对应字节数组（列优先）
    private static final byte[] PLAIN = {
            0x32, 0x43, (byte) 0xf6, (byte) 0xa8,
            (byte)0x88, 0x5a, 0x30, (byte) 0x8d,
            0x31, 0x31, (byte) 0x98, (byte) 0xa2,
            (byte) 0xe0, 0x37, 0x07, 0x34
    };
    // 密钥对应字节数组（列优先）
    private static final byte[] KEY = {
            0x2b, 0x7e, 0x15, 0x16,
            0x28, (byte) 0xae, (byte) 0xd2, (byte) 0xa6,
            (byte) 0xab, (byte) 0xf7, 0x15, (byte) 0x88,
            0x09, (byte) 0xcf, 0x4f, 0x3c
    };

    public static void main(String[] args) {
        // 密钥扩展
        byte[][] roundKeys = keyExpansion(KEY);

        // 加密过程
        byte[] encrypted = aesEncrypt(PLAIN, roundKeys);
        System.out.println("加密结果：");
        System.out.println(encrypted);

        // 解密过程
        byte[] decrypted = aesDecrypt(encrypted, roundKeys);
        System.out.println("\n解密结果：");
        System.out.println(decrypted);
    }

    // 密钥扩展算法
    private static byte[][] keyExpansion(byte[] key) {
        // 4*44的密钥扩展数组
        byte[][] w = new byte[44][4];

        // 初始密钥填充
        for (int i=0; i<4; i++) {
            System.arraycopy(key, i*4, w[i], 0, 4);
        }

        // 生成扩展密钥
        for (int i=4; i<44; i++) {
            byte[] temp = Arrays.copyOf(w[i-1], 4);

            if (i % 4 == 0) {
                // g函数处理
                rotateWord(temp);
                subWord(temp);
                temp[0] ^= RCON[i/4 -1];
            }

            // 异或运算
            for (int j=0; j<4; j++) {
                w[i][j] = (byte)(w[i-4][j] ^ temp[j]);
            }
        }
        return w;
    }

    // 密钥的循环移位
    public static void rotateWord(byte[] word) {
        byte temp = word[0];

        // 循环左移
        for (int i = 0; i < 3; i++) {
            word[i] = word[i + 1];
        }

        // 将原第一个字节放到最后一位
        word[3] = temp;
    }

    // 密钥扩展时的字替换
    private static void subWord(byte[] word) {
        for (int i = 0; i < 4; i++) {
            word[i] = (byte) S_BOX[word[i] & 0xFF];
        }
    }

    // 字节数组转为状态矩阵
    private static byte[][] bytesToState(byte[] input) {

        byte[][] state = new byte[4][4];
        for (int col = 0; col < 4; col++) {
            for (int row = 0; row < 4; row++) {
                int index = col * 4 + row;
                state[row][col] = input[index];
            }
        }
        return state;
    }

    // 状态矩阵转为字节数组
    private static byte[] stateToBytes(byte[][] state) {

        byte[] output = new byte[16];
        for (int col = 0; col < 4; col++) {
            for (int row = 0; row < 4; row++) {
                int index = col * 4 + row;
                output[index] = state[row][col];
            }
        }
        return output;
    }

    // 密钥的异或操作
    private static void addRoundKey(byte[][] state, byte[][] roundKeys, int round) {

        // 将当前轮的轮密钥转换为4x4状态矩阵
        byte[] roundKeyBytes = roundKeys[round];
        byte[][] roundKeyState = bytesToState(roundKeyBytes);

        // 执行按字节异或操作
        for (int row = 0; row < 4; row++) {
            for (int col = 0; col < 4; col++) {
                state[row][col] ^= roundKeyState[row][col];
            }
        }
    }

    // 加密主函数
    private static byte[] aesEncrypt(byte[] input, byte[][] roundKeys) {
        byte[][] state = bytesToState(input);

        // 初始轮密钥加
        addRoundKey(state, roundKeys, 0);

        // 9轮主循环
        for (int round=1; round<10; round++) {
            subBytes(state);
            shiftRows(state);
            mixColumns(state);
            addRoundKey(state, roundKeys, round);
        }

        // 最后一轮
        subBytes(state);
        shiftRows(state);
        addRoundKey(state, roundKeys, 10);

        return stateToBytes(state);
    }

    // 解密主函数
    private static byte[] aesDecrypt(byte[] input, byte[][] roundKeys) {
        byte[][] state = bytesToState(input);

        // 初始轮密钥加（使用最后轮密钥）
        addRoundKey(state, roundKeys, 10);

        // 9轮主循环
        for (int round=9; round>0; round--) {
            invShiftRows(state);
            invSubBytes(state);
            addRoundKey(state, roundKeys, round);
            invMixColumns(state);
        }

        // 最后一轮
        invShiftRows(state);
        invSubBytes(state);
        addRoundKey(state, roundKeys, 0);

        return stateToBytes(state);
    }

    // 字节替换
    private static void subBytes(byte[][] state) {
        for (int i=0; i<4; i++) {
            for (int j=0; j<4; j++) {
                int value = state[i][j] & 0xFF;
                state[i][j] = (byte)(S_BOX[value]);
            }
        }
    }

    // 行移位
    private static void shiftRows(byte[][] state) {
        // 第0行不移位
        // 第1行左移1字节
        byte temp = state[1][0];
        System.arraycopy(state[1], 1, state[1], 0, 3);
        state[1][3] = temp;

        // 第2行左移2字节（交换两次）
        swap(state[2], 0, 2);
        swap(state[2], 1, 3);

        // 第3行左移3字节（等于右移1字节）
        temp = state[3][3];
        System.arraycopy(state[3], 0, state[3], 1, 3);
        state[3][0] = temp;
    }

    // 行移位逆操作
    private static void invShiftRows(byte[][] state) {
        // 第0行不移位

        // 第1行右移1字节（等价于左移3字节）
        byte temp = state[1][3];
        state[1][3] = state[1][2];
        state[1][2] = state[1][1];
        state[1][1] = state[1][0];
        state[1][0] = temp;

        // 第2行右移2字节（交换两次）
        swap(state[2], 0, 2);
        swap(state[2], 1, 3);

        // 第3行右移3字节（等价于左移1字节）
        temp = state[3][0];
        System.arraycopy(state[3], 1, state[3], 0, 3);
        state[3][3] = temp;
    }

    // 辅助交换方法
    private static void swap(byte[] row, int i, int j) {
        byte temp = row[i];
        row[i] = row[j];
        row[j] = temp;
    }

    //
    private static final int[][] MIX_MATRIX = {
            {0x02, 0x03, 0x01, 0x01},
            {0x01, 0x02, 0x03, 0x01},
            {0x01, 0x01, 0x02, 0x03},
            {0x03, 0x01, 0x01, 0x02}
    };

    private static void mixColumns(byte[][] state) {

        // 操作状态矩阵
        for (int col = 0; col < 4; col++) {
            final byte s0 = state[0][col];
            final byte s1 = state[1][col];
            final byte s2 = state[2][col];
            final byte s3 = state[3][col];

            // 手动展开矩阵乘法运算
            state[0][col] = (byte) (
                    gfMultiply(MIX_MATRIX[0][0], s0) ^
                            gfMultiply(MIX_MATRIX[0][1], s1) ^
                            gfMultiply(MIX_MATRIX[0][2], s2) ^
                            gfMultiply(MIX_MATRIX[0][3], s3)
            );

            state[1][col] = (byte) (
                    gfMultiply(MIX_MATRIX[1][0], s0) ^
                            gfMultiply(MIX_MATRIX[1][1], s1) ^
                            gfMultiply(MIX_MATRIX[1][2], s2) ^
                            gfMultiply(MIX_MATRIX[1][3], s3)
            );

            state[2][col] = (byte) (
                    gfMultiply(MIX_MATRIX[2][0], s0) ^
                            gfMultiply(MIX_MATRIX[2][1], s1) ^
                            gfMultiply(MIX_MATRIX[2][2], s2) ^
                            gfMultiply(MIX_MATRIX[2][3], s3)
            );

            state[3][col] = (byte) (
                    gfMultiply(MIX_MATRIX[3][0], s0) ^
                            gfMultiply(MIX_MATRIX[3][1], s1) ^
                            gfMultiply(MIX_MATRIX[3][2], s2) ^
                            gfMultiply(MIX_MATRIX[3][3], s3)
            );
        }
    }

    // 逆字节替换
    private static void invSubBytes(byte[][] state) {
        for (int i = 0; i < 4; i++) {
            for (int j = 0; j < 4; j++) {
                int unsignedByte = state[i][j] & 0xFF; // 转为无符号索引
                state[i][j] = (byte) INV_S_BOX[unsignedByte];
            }
        }
    }

    // 逆列混淆
    private static void invMixColumns(byte[][] state) {
        for (int col = 0; col < 4; col++) {
            byte[] column = new byte[4];
            for (int row = 0; row < 4; row++) {
                column[row] = state[row][col];
            }

            // 逆矩阵乘法
            state[0][col] = (byte) (
                    gfMultiply(0x0E, column[0]) ^
                            gfMultiply(0x0B, column[1]) ^
                            gfMultiply(0x0D, column[2]) ^
                            gfMultiply(0x09, column[3])
            );
            state[1][col] = (byte) (
                    gfMultiply(0x09, column[0]) ^
                            gfMultiply(0x0E, column[1]) ^
                            gfMultiply(0x0B, column[2]) ^
                            gfMultiply(0x0D, column[3])
            );
            state[2][col] = (byte) (
                    gfMultiply(0x0D, column[0]) ^
                            gfMultiply(0x09, column[1]) ^
                            gfMultiply(0x0E, column[2]) ^
                            gfMultiply(0x0B, column[3])
            );
            state[3][col] = (byte) (
                    gfMultiply(0x0B, column[0]) ^
                            gfMultiply(0x0D, column[1]) ^
                            gfMultiply(0x09, column[2]) ^
                            gfMultiply(0x0E, column[3])
            );
        }
    }

    // 有限域乘法方法
    private static int gfMultiply(int a, int b) {
        int result = 0;
        for (int i = 0; i < 8; i++) {
            if ((b & 1) != 0) {
                result ^= a;
            }
            boolean carry = (a & 0x80) != 0;
            a <<= 1;
            if (carry) {
                a ^= 0x11B; // AES多项式x^8 + x^4 + x^3 + x + 1
            }
            b >>= 1;
        }
        return result;
    }
}